'''
Created on Feb 4, 2010

@author: mkiyer
'''
import logging
import os
import multiprocessing
import subprocess
import sys

import veggie.db.sample as sdb 
from veggie.db.sample.sampledb3 import Library

novoalign_quals = {'solexa-quals': 'SLXFQ',
                   'solexa1.3-quals': 'ILMFQ',
                   'phred33-quals': 'STDFQ'}
novoalign_executable = "$SW_ROOT/bioinfo/alignment/novocraft/novocraft-2.05.16/bigdaddy"
novoalign_indexes = {'hg18': 'hg18_contaminant'}
novoalign_indexes_path = "/lab/mkiyer/refdb/novocraft_indexes"

def get_indexes_path():
    return novoalign_indexes_path
def get_index(name):
    return novoalign_indexes.get(name, None)
def get_quals(fmt):
    return novoalign_quals.get(fmt,None)

def run_single_read_novoalign(library, cwd, genome):
    # skip alignments that already exist
    logging.debug("%s: preparing to run novoalign on files: %s quals: %s" % 
                  (library.id, library.fastq_files, library.fastq_format))    
    # ensure index exists
    index_name = get_index(genome)
    if index_name is None:
        raise KeyError('indexes for genome %s not available' % genome)            
    # ensure correct number of sequence files has been passed
    if len(library.fastq_files) > 2:
        raise ValueError("error in fastq_files param: cannot pass"
                         " more than 2 sequence files to novoalign")        
    # check quals format
    quals_format = library.fastq_format
    quals_param = get_quals(quals_format)        
    if quals_param == None:
        raise ValueError("invalid quality score format")

    out_file = os.path.join(cwd, library.id + '.sam')    
    err_file = os.path.join(cwd, library.id + '.log')
    # write library information
    log_file = os.path.join(cwd, library.id + '.library_info')
    log_fhd = open(log_file, 'w')
    log_fhd.write("%s\n" % library)
    log_fhd.close()
    # alignment options
    novoalign_options = ({'-o': 'SAM',
                          '-r': 'None',
                          '-d': os.path.join(novoalign_indexes_path, index_name),
                          '-f': ' '.join(library.fastq_files),
                          '-F': quals_param})
    # run the program
    arglist = [novoalign_executable]
    # option arguments
    for k, v in novoalign_options.iteritems():
        arglist.append(k)
        if v is not None:
            arglist.append(str(v))        
    argstr = ' '.join(arglist)
    # call novoalign
    #retcode = os.system('ls -l > %s 2> %s' % (out_file, err_file))
    retcode = os.system('%s > %s 2> %s' % (argstr, out_file, err_file))
    return retcode
    

def run_alignment(library, results_path, genome_version):
    if os.path.exists(results_path):
        logging.debug("Alignment results path %s exists, skipping..." % (results_path))
    else:
        # make directory for results
        logging.debug('Creating alignment results path %s' % (results_path))
        os.makedirs(results_path)
    run_single_read_novoalign(library, results_path, genome_version)

def alignment_worker(arg):
    library_analysis, alignment_path, genome_version = arg
    run_alignment(library_analysis, alignment_path, genome_version)
    return library_analysis.id

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    # TODO: add path to sampledb instance to config file
    sdbi_path = '/lab/mkiyer/refdb/sample/sampledb_2010_Feb_04.pickle'
    # load the current sampledb snapshot
    sdbi = sdb.get_sampledb_instance(sdbi_path)
    # chipseq alignment results path
    alignment_path = '/home/mkiyer/alignment'
    # genome version
    genome_version = 'hg18'
    # number of processors to run
    processors = 3
    # find all chipseq samples
    alignment_tasks = []
    for library_id, library_analyses in sdbi.libraries.iteritems():
        samples = sdbi.get_library_samples(library_id)
        assert len(samples) == 1
        sample = samples[0]        
        if sample.app_type == 'Chipseq':
            library_analysis = Library.choose_best_library(library_analyses)
            # BEGIN HACK
            old_library_id = '_'.join(library_analysis.id.split('_')[1:])
            old_library_path = os.path.join('/home/mkiyer/exper', old_library_id)
            if os.path.exists(old_library_path):
                logging.warning("Old library alignment exists at %s, skipping to next library" % (old_library_path))
                continue
            # END HACK            
            alignment_tasks.append((library_analysis, alignment_path, genome_version))
    # start pool of processes
    pool = multiprocessing.Pool(processors)
    imap_unordered_it = pool.imap_unordered(alignment_worker,
                                            alignment_tasks)
    for finished_id in imap_unordered_it:
        logging.debug("%s finished" % finished_id)
    pool.close()
    pool.join()
